#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# Define port constants
SWITCH_TO_HOST_PORT = 1  # Hosts are always connected to port 1

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr):
    """
    Installs three rules for tunnel communication:
    1) Ingress rule on ingress switch
    2) Transit rule on ingress switch
    3) Egress rule on egress switch
    """
    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed ingress tunnel rule on {ingress_sw.name}")

    # 2) Tunnel Transit Rule
    # Determine the correct output port based on switch and tunnel ID
    if ingress_sw.name == 's1':
        if tunnel_id in [100, 101]:  # s1 <-> s2 traffic
            SWITCH_TO_SWITCH_PORT = 2
        elif tunnel_id in [200, 201]:  # s1 <-> s3 traffic
            SWITCH_TO_SWITCH_PORT = 3
    elif ingress_sw.name == 's2':
        if tunnel_id in [100, 101]:  # s2 <-> s1 traffic
            SWITCH_TO_SWITCH_PORT = 2
        elif tunnel_id in [300, 301]:  # s2 <-> s3 traffic
            SWITCH_TO_SWITCH_PORT = 3
    elif ingress_sw.name == 's3':
        if tunnel_id in [200, 201]:  # s3 <-> s1 traffic
            SWITCH_TO_SWITCH_PORT = 2
        elif tunnel_id in [300, 301]:  # s3 <-> s2 traffic
            SWITCH_TO_SWITCH_PORT = 3

    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": SWITCH_TO_SWITCH_PORT
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed transit tunnel rule on {ingress_sw.name}")

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print(f"Installed egress tunnel rule on {egress_sw.name}")

def readTableRules(p4info_helper, sw):
    """Reads and prints all table entries from the switch"""
    print(f'\n----- Reading tables rules for {sw.name} -----')
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print(f'{table_name}: ', end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print(f'{p4info_helper.get_match_field_value(m)}', end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print(f'{p.value}', end=' ')
            print()

def printCounter(p4info_helper, sw, counter_name, index):
    """Prints the specified counter from the switch"""
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print(f"{sw.name} {counter_name} {index}: {counter.data.packet_count} packets ({counter.data.byte_count} bytes)")

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connections for all 3 switches
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install P4 program on all switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                     bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                     bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                     bmv2_json_file_path=bmv2_file_path)

        # Write tunnel rules for all possible paths
        # h1 <-> h2 (via s1 and s2)
        writeTunnelRules(p4info_helper, s1, s2, 100, "08:00:00:00:02:22", "10.0.2.2")
        writeTunnelRules(p4info_helper, s2, s1, 101, "08:00:00:00:01:11", "10.0.1.1")

        # h1 <-> h3 (via s1 and s3)
        writeTunnelRules(p4info_helper, s1, s3, 200, "08:00:00:00:03:33", "10.0.3.3")
        writeTunnelRules(p4info_helper, s3, s1, 201, "08:00:00:00:01:11", "10.0.1.1")

        # h2 <-> h3 (via s2 and s3)
        writeTunnelRules(p4info_helper, s2, s3, 300, "08:00:00:00:03:33", "10.0.3.3")
        writeTunnelRules(p4info_helper, s3, s2, 301, "08:00:00:00:02:22", "10.0.2.2")

        # Read and print table rules
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        # Monitor counters
        while True:
            sleep(2)
            print('\n----- Reading tunnel counters -----')
            # h1 <-> h2 counters
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)

            # h1 <-> h3 counters
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)

            # h2 <-> h3 counters
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                      type=str, action="store", required=False,
                      default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                      type=str, action="store", required=False,
                      default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)
    main(args.p4info, args.bmv2_json)